#!/usr/bin/env python3
"""
音频处理模块
提供音频降噪、修复等核心功能
"""

import os
import numpy as np
import librosa
import soundfile as sf
from typing import Optional, Tuple
import logging

# 配置日志
logger = logging.getLogger(__name__)


class AudioProcessor:
    """音频处理器类"""
    
    def __init__(self):
        """初始化音频处理器"""
        self.sample_rate = 22050  # 默认采样率
        self.auto_noise_reduction = 0.5  # 自动模式降噪强度
        
    def load_audio(self, file_path: str) -> Tuple[np.ndarray, int]:
        """
        加载音频文件
        
        Args:
            file_path: 音频文件路径
            
        Returns:
            (音频数据, 采样率)
        """
        try:
            # 使用librosa加载音频
            audio, sr = librosa.load(file_path, sr=self.sample_rate)
            logger.info(f"音频加载成功: {file_path}, 采样率: {sr}, 长度: {len(audio)}")
            return audio, sr
        except Exception as e:
            logger.error(f"音频加载失败: {file_path}, 错误: {str(e)}")
            raise
    
    def save_audio(self, audio: np.ndarray, file_path: str, sample_rate: int) -> bool:
        """
        保存音频文件
        
        Args:
            audio: 音频数据
            file_path: 输出文件路径
            sample_rate: 采样率
            
        Returns:
            是否保存成功
        """
        try:
            # 确保音频数据在有效范围内
            audio = np.clip(audio, -1.0, 1.0)
            
            # 使用soundfile保存音频
            sf.write(file_path, audio, sample_rate)
            logger.info(f"音频保存成功: {file_path}")
            return True
        except Exception as e:
            logger.error(f"音频保存失败: {file_path}, 错误: {str(e)}")
            return False
    
    def detect_noise_profile(self, audio: np.ndarray, sample_rate: int) -> np.ndarray:
        """
        检测音频中的噪声特征
        
        Args:
            audio: 音频数据
            sample_rate: 采样率
            
        Returns:
            噪声特征谱
        """
        try:
            # 使用音频的前几秒作为噪声样本（假设开头是静音或背景噪声）
            noise_duration = min(3.0, len(audio) / sample_rate)  # 最多3秒
            noise_samples = int(noise_duration * sample_rate)
            noise_audio = audio[:noise_samples]
            
            # 计算噪声的频谱特征
            noise_stft = librosa.stft(noise_audio)
            noise_profile = np.abs(noise_stft).mean(axis=1)
            
            logger.info(f"噪声特征检测完成，噪声时长: {noise_duration}秒")
            return noise_profile
            
        except Exception as e:
            logger.error(f"噪声特征检测失败: {str(e)}")
            # 返回默认噪声特征
            return np.ones(1025) * 0.1
    
    def reduce_noise(self, audio: np.ndarray, noise_profile: np.ndarray, 
                    reduction_strength: float = 0.5) -> np.ndarray:
        """
        降噪处理
        
        Args:
            audio: 原始音频数据
            noise_profile: 噪声特征
            reduction_strength: 降噪强度 (0.0-1.0)
            
        Returns:
            降噪后的音频数据
        """
        try:
            # 计算音频的STFT
            stft = librosa.stft(audio)
            magnitude = np.abs(stft)
            phase = np.angle(stft)
            
            # 应用噪声门限
            threshold = noise_profile[:, np.newaxis] * (1 + reduction_strength)
            mask = magnitude > threshold
            
            # 软掩码处理，避免过度降噪
            soft_mask = mask.astype(float) * (1 - reduction_strength * 0.5) + \
                       (~mask).astype(float) * reduction_strength * 0.5
            
            # 应用掩码
            reduced_magnitude = magnitude * soft_mask
            
            # 重构STFT
            reduced_stft = reduced_magnitude * np.exp(1j * phase)
            reduced_audio = librosa.istft(reduced_stft)
            
            # 确保长度一致
            if len(reduced_audio) < len(audio):
                reduced_audio = np.pad(reduced_audio, (0, len(audio) - len(reduced_audio)))
            else:
                reduced_audio = reduced_audio[:len(audio)]
            
            logger.info(f"降噪处理完成，强度: {reduction_strength}")
            return reduced_audio
            
        except Exception as e:
            logger.error(f"降噪处理失败: {str(e)}")
            return audio
    
    def remove_clicks_pops(self, audio: np.ndarray, sample_rate: int,
                           threshold: float = 0.8) -> np.ndarray:
        """
        去除爆音和点击声
        
        Args:
            audio: 音频数据
            sample_rate: 采样率
            threshold: 检测阈值
            
        Returns:
            处理后的音频数据
        """
        try:
            # 检测音频中的突变点
            audio_diff = np.diff(audio)
            abs_diff = np.abs(audio_diff)
            
            # 设置动态阈值
            median_diff = np.median(abs_diff)
            dynamic_threshold = median_diff * 10  # 经验值
            
            # 找到异常点
            click_indices = np.where(abs_diff > dynamic_threshold)[0]
            
            if len(click_indices) > 0:
                logger.info(f"检测到 {len(click_indices)} 个爆音点")
                
                # 使用插值修复异常点
                repaired_audio = audio.copy()
                
                for idx in click_indices:
                    # 获取周围样本
                    start = max(0, idx - 10)
                    end = min(len(audio), idx + 11)
                    
                    # 线性插值
                    if start > 0 and end < len(audio):
                        start_val = audio[start]
                        end_val = audio[end - 1]
                        num_samples = end - start
                        
                        interpolation = np.linspace(start_val, end_val, num_samples)
                        repaired_audio[start:end] = interpolation
                
                return repaired_audio
            else:
                logger.info("未检测到爆音")
                return audio
                
        except Exception as e:
            logger.error(f"爆音去除失败: {str(e)}")
            return audio
    
    def reduce_hum_noise(self, audio: np.ndarray, sample_rate: int,
                        freq_range: Tuple[float, float] = (50, 60)) -> np.ndarray:
        """
        减少电流声（嗡嗡声）
        
        Args:
            audio: 音频数据
            sample_rate: 采样率
            freq_range: 电流声频率范围 (默认50-60Hz)
            
        Returns:
            处理后的音频数据
        """
        try:
            # 设计带阻滤波器
            from scipy import signal
            
            # 转换为Hz
            nyquist = sample_rate / 2
            low_freq = freq_range[0] / nyquist
            high_freq = freq_range[1] / nyquist
            
            # 设计带阻滤波器
            b, a = signal.butter(4, [low_freq, high_freq], btype='bandstop')
            
            # 应用滤波器
            filtered_audio = signal.filtfilt(b, a, audio)
            
            logger.info(f"电流声降噪完成，频率范围: {freq_range}Hz")
            return filtered_audio
            
        except Exception as e:
            logger.error(f"电流声降噪失败: {str(e)}")
            return audio
    
    def normalize_audio(self, audio: np.ndarray, target_db: float = -3.0) -> np.ndarray:
        """
        音频归一化
        
        Args:
            audio: 音频数据
            target_db: 目标响度 (dB)
            
        Returns:
